//
// Created by hjhu on 2024/7/2.
//

#include "Fan.h"
#include "HourGlass.h"


namespace coastal {

    FanImpl::FanImpl(int64_t
                     num_modules
    )
            :
            num_modules(num_modules) {
        conv1 = torch::nn::Conv2d(torch::nn::Conv2dOptions(3, 64, 7).stride(2).padding(3));
        bn1 = torch::nn::BatchNorm2d(64);
        conv2 = std::make_shared<coastal::ConvBlockImpl>(64, 128);
        conv3 = std::make_shared<coastal::ConvBlockImpl>(128, 128);
        conv4 = std::make_shared<coastal::ConvBlockImpl>(128, 256);

        for (int64_t hg_module = 0; hg_module < num_modules; ++hg_module) {
            register_module("m" + std::to_string(hg_module), std::make_shared<coastal::HourGlassImpl>(1, 4, 256));
            register_module("top_m_" + std::to_string(hg_module), std::make_shared<coastal::ConvBlockImpl>(256, 256));
            register_module("conv_last" + std::to_string(hg_module),
                            torch::nn::Conv2d(torch::nn::Conv2dOptions(256, 256, 1).stride(1).padding(0)));
            register_module("bn_end" + std::to_string(hg_module), nn::BatchNorm2d(256));
            register_module("l" + std::to_string(hg_module),
                            torch::nn::Conv2d(torch::nn::Conv2dOptions(256, 68, 1).stride(1).padding(0)));

            if (hg_module < num_modules - 1) {
                register_module("bl" + std::to_string(hg_module),
                                torch::nn::Conv2d(torch::nn::Conv2dOptions(256, 256, 1).stride(1).padding(0)));
                register_module("al" + std::to_string(hg_module),
                                torch::nn::Conv2d(torch::nn::Conv2dOptions(68, 256, 1).stride(1).padding(0)));
            }
        }
    }

    torch::Tensor FanImpl::forward(torch::Tensor x) {
        x = torch::relu(bn1->forward(conv1->forward(x)));
        x = torch::avg_pool2d(conv2->forward(x), 2, 2);
        x = conv3->forward(x);
        x = conv4->forward(x);

        auto previous = x;
        std::vector<torch::Tensor> outputs;

        for (int i = 0; i < num_modules; ++i) {
            auto hg = get_module<coastal::ConvBlockImpl>("m" + std::to_string(i))->forward(previous);
            auto ll = get_module<coastal::ConvBlockImpl>("top_m_" + std::to_string(i))->forward(hg);
            ll = torch::relu(get_module<torch::nn::Conv2dImpl>("bn_end" + std::to_string(i))->forward(
                    get_module<torch::nn::Conv2d>("conv_last" + std::to_string(i))->forward(ll)));

            auto tmp_out = get_module<torch::nn::Conv2dImpl>("l" + std::to_string(i))->forward(ll);
            outputs.push_back(tmp_out);
            if (i < num_modules - 1) {
                ll = get_module<torch::nn::Conv2dImpl>("bl" + std::to_string(i))->forward(ll);
                auto tmp_out_ = get_module<torch::nn::Conv2dImpl>("al" + std::to_string(i))->forward(tmp_out);
                previous = previous + ll + tmp_out_;
            }
        }

        return torch::stack(outputs);
    }


} // coastal